package bitset

import (
	"bytes"
	"fmt"
)

const (
	mask  = 0x3F // 63的二进制表示
	shift = 6    // 2^6 = 64
)

// BitSet 定义位集合的数据结构
type BitSet struct {
	data []uint64
	size int
}

// New 创建一个新的BitSet，初始容量为n位
func New(n int) *BitSet {
	if n < 0 {
		n = 0
	}
	return &BitSet{
		data: make([]uint64, (n+63)/64),
		size: n,
	}
}

// Set 将指定位置的位设置为1
func (b *BitSet) Set(pos int) *BitSet {
	if pos < 0 || pos >= b.size {
		return b
	}
	idx := pos >> shift
	b.data[idx] |= 1 << (pos & mask)
	return b
}

// Clear 将指定位置的位设置为0
func (b *BitSet) Clear(pos int) *BitSet {
	if pos < 0 || pos >= b.size {
		return b
	}
	idx := pos >> shift
	b.data[idx] &^= 1 << (pos & mask)
	return b
}

// Test 检查指定位置的位是否为1
func (b *BitSet) Test(pos int) bool {
	if pos < 0 || pos >= b.size {
		return false
	}
	idx := pos >> shift
	return b.data[idx]&(1<<(pos&mask)) != 0
}

// Resize 调整BitSet的大小
func (b *BitSet) Resize(newSize int) *BitSet {
	if newSize < 0 {
		newSize = 0
	}

	newLen := (newSize + 63) / 64
	if newLen > len(b.data) {
		newData := make([]uint64, newLen)
		copy(newData, b.data)
		b.data = newData
	} else if newLen < len(b.data) {
		// 清除超出新大小范围的高位
		lastIdx := newLen - 1
		if lastIdx >= 0 {
			// 计算最后一个uint64中有效的位数
			validBits := newSize % 64
			if validBits == 0 {
				validBits = 64
			}
			// 清除超出有效位数的高位
			mask := uint64((1 << validBits) - 1)
			b.data[lastIdx] &= mask
		}
		b.data = b.data[:newLen]
	}
	b.size = newSize
	return b
}

// Len 返回BitSet的位数
func (b *BitSet) Len() int {
	return b.size
}

// Count 返回BitSet中1的数量
func (b *BitSet) Count() int {
	count := 0
	for _, v := range b.data {
		count += popCount(v)
	}
	return count
}

// 计算一个uint64中1的个数（使用SWAR算法）
func popCount(x uint64) int {
	x = x - ((x >> 1) & 0x5555555555555555)
	x = (x & 0x3333333333333333) + ((x >> 2) & 0x3333333333333333)
	x = (x + (x >> 4)) & 0x0f0f0f0f0f0f0f0f
	x = x + (x >> 8)
	x = x + (x >> 16)
	x = x + (x >> 32)
	return int(x & 0x7f)
}

// Union 计算当前BitSet与另一个BitSet的并集
func (b *BitSet) Union(other *BitSet) *BitSet {
	minLen := min(len(b.data), len(other.data))
	for i := 0; i < minLen; i++ {
		b.data[i] |= other.data[i]
	}
	return b
}

// Intersect 计算当前BitSet与另一个BitSet的交集
func (b *BitSet) Intersect(other *BitSet) *BitSet {
	minLen := min(len(b.data), len(other.data))
	for i := 0; i < minLen; i++ {
		b.data[i] &= other.data[i]
	}
	// 清除超出other长度的位
	for i := minLen; i < len(b.data); i++ {
		b.data[i] = 0
	}
	return b
}

// Difference 计算当前BitSet与另一个BitSet的差集
func (b *BitSet) Difference(other *BitSet) *BitSet {
	minLen := min(len(b.data), len(other.data))
	for i := 0; i < minLen; i++ {
		b.data[i] &^= other.data[i]
	}
	return b
}

// Complement 计算BitSet的补集
func (b *BitSet) Complement() *BitSet {
	for i := range b.data {
		b.data[i] = ^b.data[i]
	}
	// 清除超出size的高位
	if b.size > 0 {
		lastIdx := (b.size - 1) >> shift
		validBits := b.size % 64
		if validBits == 0 {
			validBits = 64
		}
		mask := uint64((1 << validBits) - 1)
		b.data[lastIdx] &= mask
	}
	return b
}

// Clone 创建BitSet的副本
func (b *BitSet) Clone() *BitSet {
	dataCopy := make([]uint64, len(b.data))
	copy(dataCopy, b.data)
	return &BitSet{
		data: dataCopy,
		size: b.size,
	}
}

// Equal 检查两个BitSet是否相等
func (b *BitSet) Equal(other *BitSet) bool {
	if b.size != other.size {
		return false
	}

	minLen := min(len(b.data), len(other.data))
	for i := 0; i < minLen; i++ {
		if b.data[i] != other.data[i] {
			return false
		}
	}

	// 检查较长BitSet的剩余部分是否全为0
	if len(b.data) > minLen {
		for i := minLen; i < len(b.data); i++ {
			if b.data[i] != 0 {
				return false
			}
		}
	} else if len(other.data) > minLen {
		for i := minLen; i < len(other.data); i++ {
			if other.data[i] != 0 {
				return false
			}
		}
	}

	return true
}

// String 返回BitSet的字符串表示形式
func (b *BitSet) String() string {
	var buf bytes.Buffer
	buf.WriteByte('[')
	for i := 0; i < b.size; i++ {
		if b.Test(i) {
			fmt.Fprintf(&buf, "%d ", i)
		}
	}
	if buf.Len() > 1 {
		buf.Truncate(buf.Len() - 1) // 移除最后一个空格
	}
	buf.WriteByte(']')
	return buf.String()
}

// ForEach 遍历所有设置为1的位，并对每个位执行回调函数
func (b *BitSet) ForEach(fn func(pos int) bool) {
	for i, v := range b.data {
		if v == 0 {
			continue
		}
		base := i << shift
		for j := 0; j < 64; j++ {
			if v&(1<<j) != 0 {
				pos := base + j
				if pos >= b.size {
					break
				}
				if !fn(pos) {
					return
				}
			}
		}
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
